package fun.ticsmyc.multiThread.流控;

import java.util.concurrent.atomic.AtomicLong;

public class 令牌桶 {


    //桶容量
    private final long capacity;

    //令牌填充速率
    private final long refillRate;

    //当前令牌数量
    private AtomicLong tokens;

    //上次填充时间戳
    private long lastRefillTime;

    public 令牌桶(long capacity, long refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = new AtomicLong(capacity);
        this.lastRefillTime = System.nanoTime();
    }

    public synchronized boolean tryAcquire(){
        refillTokens();
        if(tokens.get()>0){
            tokens.decrementAndGet();
            return true;
        }
        return false;
    }
    public synchronized boolean tryAcquire(int numTokens){
        refillTokens();
        if(tokens.get() >= numTokens){
            tokens.addAndGet(-numTokens);
            return true;
        }
        return false;
    }


    private synchronized void refillTokens() {
        long now = System.nanoTime();

        //间隔xx秒
        double elapsedSec = (now-this.lastRefillTime)*1e-9;

        long tokensToAdd = (long)(elapsedSec * refillRate);

        if(tokensToAdd >0){
            tokens.set(Math.min(capacity, tokensToAdd+tokens.get()));
            lastRefillTime = now;
        }

    }

    public static void main(String[] args) throws InterruptedException {
        令牌桶 bucket = new 令牌桶(10,2);

        for (int i = 0; i < 50; i++) {
            if(bucket.tryAcquire()){
                System.out.println("请求"+i+"通过");
            }else{
                System.out.println("请求"+i+" 限流");
            }
            Thread.sleep(100);
        }

    }
}
